
import os
import aiohttp
import logging
import openai

logger = logging.getLogger("ai.engine")

class AIEngine:
    def __init__(self, config):
        self.config = config
        self.ai_conf = config.get("ai", {})
        self.deepseek_conf = config.get("deepseek", {})
        self.xiaoi_conf = config.get("xiaoi", {})
        self.provider = self.ai_conf.get("provider", "openai").lower()
        self.use_deepseek = self.provider == "deepseek"
        self.use_xiaoi = self.provider == "xiaoi"
        if self.use_deepseek:
            base_url = self.deepseek_conf.get("base_url")
            if base_url:
                if base_url.endswith("/beta"):
                    endpoint = base_url.rstrip("/") + "/chat/completions"
                else:
                    endpoint = base_url.rstrip("/") + "/v1/chat/completions"
                self.deepseek_conf["endpoint"] = endpoint
            elif not self.deepseek_conf.get("endpoint"):
                self.deepseek_conf["endpoint"] = "https://api.deepseek.com/v1/chat/completions"
            logger.info(f"AI引擎初始化，模型: {self.deepseek_conf.get('model')}, endpoint: {self.deepseek_conf.get('endpoint')}")
        elif self.use_xiaoi:
            logger.info(f"AI引擎初始化，小爱AI，模型: {self.xiaoi_conf.get('model')}, url: {self.xiaoi_conf.get('api_url')}")
        else:
            openai.api_key = self.ai_conf.get("api_key")
            logger.info(f"AI引擎初始化，模型: {self.ai_conf.get('model')}")

    async def get_ai_response(self, prompt):
        ai_danmu_limit = self.config.get("ai_danmu_limit", {})
        concise_mode = ai_danmu_limit.get("concise_mode", False)
        max_length = ai_danmu_limit.get("max_length", 30)
        system_prompt = self.ai_conf.get("system_prompt", "")
        if concise_mode:
            system_prompt = system_prompt + f"（请用极简短的语气回复，务必控制在{max_length}汉字或60字符以内，超出将被丢弃）"
        if self.use_deepseek:
            return await self._get_deepseek_response(prompt, system_prompt=system_prompt)
        elif self.use_xiaoi:
            return await self._get_xiaoi_response(prompt, system_prompt=system_prompt)
        else:
            return await self._get_openai_response(prompt, system_prompt=system_prompt)

    async def _get_openai_response(self, prompt, system_prompt=None):
        try:
            response = await openai.ChatCompletion.acreate(
                model=self.ai_conf.get("model", "gpt-3.5-turbo"),
                messages=[
                    {"role": "system", "content": system_prompt if system_prompt is not None else self.ai_conf.get("system_prompt", "")},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=self.ai_conf.get("max_response_length", 100),
                temperature=self.ai_conf.get("temperature", 0.7)
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            logger.error(f"OpenAI请求失败: {e}")
            return "AI服务暂时不可用，请稍后再试。"

    async def _get_deepseek_response(self, prompt, system_prompt=None):
        try:
            headers = {
                "Authorization": f"Bearer {self.deepseek_conf.get('api_key')}",
                "Content-Type": "application/json"
            }
            max_tokens = self.deepseek_conf.get("max_tokens", 100)
            try:
                max_tokens = int(max_tokens)
            except Exception:
                max_tokens = 100
            temperature = self.deepseek_conf.get("temperature", 0.7)
            try:
                temperature = float(temperature)
            except Exception:
                temperature = 0.7
            payload = {
                "model": self.deepseek_conf.get("model", "deepseek-chat"),
                "messages": [
                    {"role": "system", "content": system_prompt if system_prompt is not None else self.ai_conf.get("system_prompt", "")},
                    {"role": "user", "content": prompt}
                ],
                "max_tokens": max_tokens,
                "temperature": temperature
            }
            async with aiohttp.ClientSession() as session:
                async with session.post(self.deepseek_conf.get("endpoint"), json=payload, headers=headers, timeout=30) as resp:
                    data = await resp.json()
                    return data["choices"][0]["message"]["content"].strip()
        except Exception as e:
            logger.error(f"DeepSeek请求失败: {e}")
            return "AI服务暂时不可用，请稍后再试。"

    async def _get_xiaoi_response(self, prompt, system_prompt=None):
        try:
            headers = {
                "Authorization": f"Bearer {self.xiaoi_conf.get('api_key')}",
                "Content-Type": "application/json"
            }
            temperature = self.xiaoi_conf.get("temperature", 0.7)
            try:
                temperature = float(temperature)
            except Exception:
                temperature = 0.7
            payload = {
                "model": self.xiaoi_conf.get("model", "xiaoi-chat"),
                "messages": [
                    {"role": "system", "content": system_prompt if system_prompt is not None else self.ai_conf.get("system_prompt", "")},
                    {"role": "user", "content": prompt}
                ],
                "temperature": temperature
            }
            async with aiohttp.ClientSession() as session:
                async with session.post(self.xiaoi_conf.get("api_url"), json=payload, headers=headers, timeout=30) as resp:
                    data = await resp.json()
                    # 假设返回格式与openai兼容
                    return data["choices"][0]["message"]["content"].strip()
        except Exception as e:
            logger.error(f"小爱AI请求失败: {e}")
            return "AI服务暂时不可用，请稍后再试。"